[atom-vllm]: enable DSpark speculative decoding for DeepSeek-V4-Flash-0731 - #2137

Motivation

Closes #2136.

DeepSeek-V4-Flash-0731 runs with dspark in ATOM native mode but not through the vLLM plugin. This PR brings the plugin to parity: DSpark as a draft method via --speculative-config '{"method":"dspark"}', plus four correctness fixes on the existing V4 plugin path that block the model regardless of spec decode.

Technical Details

Model. DeepseekV4Model.forward optionally returns per-layer aux hidden states for EAGLE3/DSpark/DFlash drafts, averaged over the mHC dim; DeepseekV4ForCausalLM PCP-all-gathers them outside the compile boundary. The aux branch sits inside @support_torch_compile, so aux_hidden_state_layers is baked into the traced graph and must be set at load time: from --level 2 up, the dispatcher replays code object 0 without evaluating guards, so later mutation is silently ignored. Noted in a COMPILE BOUNDARY comment in atom/models/deepseek_v4.py.

KV cache. The draft needs a second KV group (sliding-window MLA, block 64) beside the V4 proxy's block 128; dspark_draft_kv_patch registers it lazily and exposes it for pickle via PEP-562 __getattr__. V4's SWA ring is not keyed by a vLLM block, so a prefix-cache hit reads a stale window; every hit is rolled back by max(win_with_spec, index_topk) tokens and the tail re-forwarded, converted per group since the block sizes differ.

KNOWN ISSUE: the index_topk term is empirical (512 on V4-Flash vs a 128-token window), and the same floor reproduces with prefix caching off, so the real defect is in the sparse indexer. This rollback is a workaround; happy to track it separately.

Correctness fixes (independent of DSpark, latent on the existing V4 path):

• VocabParallelEmbedding.forward used a raw F.embedding on the tp_size == 1 branch, which GPU-faults on the -1 spec placeholder async scheduling can emit. Now replicated_embedding, which already masks it; its comment says VocabParallelEmbedding routes through the masked op, true only of the TP>1 branch.
• Async scheduling drops the sampled-token sync, so pinned staging buffers can tear; gated with a depth-1 torch.cuda.Event.
• Idle DP ranks reach build() via execute_dummy_batch() with seq_len == query_len; synthesized a decode context.
• Piecewise cudagraph modes are demoted on the plugin path, a demotion vLLM skips when it compiles nothing.

Perf. Request ids reach ATOM metadata builders through a thread-local instead of a per-step block_table[:, 0] D2H copy.

Test Plan

New: test_vllm_dspark_draft_kv_patch.py, test_vllm_v4_async_scheduling_guards.py. Extended: test_vllm_026_compat.py, test_vllm_kimi_k3.py.

# 3 sglang files cannot be collected here (no sglang module), which aborts the whole dir
python -m pytest tests/plugin -q \
  --ignore=tests/plugin/test_sglang_gdn_target_verify.py \
  --ignore=tests/plugin/test_sglang_gdn_verify_batched_ssm.py \
  --ignore=tests/plugin/test_sglang_model_arch_metadata.py
black --check . && ruff check .

Run on this branch (6fef3b9d) and its base (412f5bfe) to separate regressions from pre-existing failures. Accuracy: DeepSeek-V4-Flash-0731, 8x MI355X (gfx950), TP8, GSM8K 5-shot, 200 questions, via lm_eval local-completions against a live server, with a spec-off control and a third arm at temperature 1.0 to exercise the stochastic rejection path.

Test Result

base 412f5bfe : 37 failed, 174 passed
this branch   : 36 failed, 209 passed
new failures  : none

All 36 are pre-existing and unrelated (27 sglang, 9 rtpllm), failing identically on base. black --check . clean; ruff on the touched files is 24 findings vs 33 at base.

GSM8K exact_match (strict):

DSpark costs 0.005 against the control, inside one stderr, and both greedy arms clear the 0.94 bar the nightly DeepSeek-V4-Pro TP8 case uses. The temperature arm differs from greedy on 197/200 questions, confirming the stochastic verifier is engaged: a draft that ignored the requested distribution would return the greedy text and score identically.

Prefix-cache sweep on V4-Flash-0731: rollback 128/256/384 -> 1/6 correct, 512 -> 6/6, and 512 holds at 2K/6K/17K prefixes.

Not run: throughput benchmarks. .github/workflows/atom-vllm-test.yaml is gated behind ci:full or an approving review, so it needs a maintainer to trigger. No CI case passes --speculative-config today, so none exercises dspark; glad to add one.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.